import matplotlib.pyplot as plt
import networkx as nx
import matplotlib.animation as animation

# 定义一个简单的加权无向图
edges = [
    ('A', 'B', 4),
    ('A', 'H', 8),
    ('B', 'H', 11),
    ('B', 'C', 8),
    ('H', 'I', 7),
    ('H', 'G', 1),
    ('I', 'G', 6),
    ('C', 'I', 2),
    ('C', 'D', 7),
    ('C', 'F', 4),
    ('G', 'F', 2),
    ('D', 'E', 9),
    ('D', 'F', 14),
    ('F', 'E', 10),
]

# 创建加权无向图
G = nx.Graph()
G.add_weighted_edges_from(edges)

# 计算 Prim's 算法生成的最小生成树
mst_edges = list(nx.minimum_spanning_edges(G, algorithm='prim', data=False))

# 动画展示 Prim's 算法生成最小生成树的过程
fig, ax = plt.subplots(figsize=(8, 6))

def update(num):
    ax.clear()
    pos = nx.spring_layout(G)

    # 绘制节点和边
    nx.draw_networkx_nodes(G, pos, ax=ax, node_color='lightblue', node_size=500)
    nx.draw_networkx_edges(G, pos, ax=ax, edge_color='gray')
    nx.draw_networkx_labels(G, pos, ax=ax, font_color='black')

    # 绘制当前步骤的最小生成树边
    current_edges = mst_edges[:num+1]
    nx.draw_networkx_edges(G, pos, edgelist=current_edges, edge_color='red', ax=ax)

    ax.set_title(f'Prim\'s Algorithm Step {num+1}')

ani = animation.FuncAnimation(fig, update, frames=len(mst_edges), interval=2000, repeat=False)

plt.show()
